# -*-coding:utf-8 -*-
'''
python demo_avi.py configs/tusimple.py --test_model ./weights/tusimple_18.pth
python demo_avi.py configs/tusimple.py --test_model D:\code\Fast-Lane-Detection\logs\20221026_164137_lr_4e-04_b_32\ep049.pth

# 只使用cpu就25行：`device = 'cpu'`
# 根据自己数据类别调整 `num_lanes = 2`
'''
import cv2
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print
import torch
import scipy.special
import numpy as np
import torchvision.transforms as transforms
from data.constant import culane_row_anchor, tusimple_row_anchor
from PIL import Image
import time

if __name__ == "__main__":
    # 是否使用cpu
    print("cuda:{}", torch.cuda.is_available())
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # 只使用cpu就`device = 'cpu'`
    device = 'cpu'
    num_lanes = 2 # 根据自己数据类别调整
    # 读取数据
    cap = cv2.VideoCapture("test.mp4")

# ----------------------------------

    # 保存视频的类
    # fourcc = cv2.VideoWriter_fourcc(*'MJPG')
    # vout = cv2.VideoWriter(str(111) + '.avi', fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

    print("w = {},h = {}".format(cap.get(3), cap.get(4)))

    # 加载参数
    torch.backends.cudnn.benchmark = True # 那么cuDNN使用的非确定性算法就会自动寻找最适合当前配置的高效算法，来达到优化运行效率的问题
    args, cfg = merge_config()
    dist_print('start testing...')
    assert cfg.backbone in ['18', '34', '50', '101', '152', '50next', '101next', '50wide', '101wide']

    img_w, img_h = 1280, 720
    if cfg.dataset == 'CULane':
        cls_num_per_lane = 18
        row_anchor = culane_row_anchor
    elif cfg.dataset == 'Tusimple':
        cls_num_per_lane = 56
        row_anchor = tusimple_row_anchor
    else:
        raise NotImplementedError

    net = parsingNet(pretrained=False, backbone=cfg.backbone, cls_dim=(cfg.griding_num + 1, cls_num_per_lane, num_lanes),
                     use_aux=False).to(device)  # we dont need auxiliary segmentation in testing

    state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
    compatible_state_dict = {}

    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()

    # 图像变换
    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])


    while True:
        rval, frame = cap.read()
        frame = cv2.resize(frame, (img_w, img_h))
        if rval == False:
            print("Image source is error...")
            break
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img_ = Image.fromarray(img)  # 实现array到image的转换
        imgs = img_transforms(img_)
        imgs = imgs.unsqueeze(0)  # 起到升维的作用

        imgs = imgs.to(device)

        t1 = time.time()

        with torch.no_grad():
            out = net(imgs)

        t2 = time.time()

        col_sample = np.linspace(0, 800 - 1, cfg.griding_num)
        col_sample_w = col_sample[1] - col_sample[0]

        out_j = out[0].data.cpu().numpy()
        out_j = out_j[:, ::-1, :]
        prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
        idx = np.arange(cfg.griding_num) + 1
        idx = idx.reshape(-1, 1, 1)
        loc = np.sum(prob * idx, axis=0)
        out_j = np.argmax(out_j, axis=0)
        loc[out_j == cfg.griding_num] = 0
        out_j = loc

        t3 = time.time()

        # print((t2 - t1))
        print('Done.({:.2}s)'.format(t2 - t1))

        for i in range(out_j.shape[1]):
            if np.sum(out_j[:, i] != 0) > 2:
                # print("--------------")
                points = []
                for k in range(out_j.shape[0]):
                    if out_j[k, i] > 0:
                        ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1, int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)  # 恢复尺度

                        # cv2.circle(frame, ppp, 3, (0, 255, 0), -1)
                        # p1 = (int(out_j[k + 1, i] * col_sample_w * img_w / 800) - 1,
                        #       int(img_h * (row_anchor[cls_num_per_lane - 1 - (k + 1)] / 288)) - 1)
                        # if (p1[0] > 0 and p1[1] > 0):
                        #     cv2.line(frame, ppp, p1, (255, 0, 0), 2)

                # ----------------统计了画线段----------------------------
                        points.append(ppp)
                for i, p in enumerate(points):
                    cv2.circle(frame, p, 3, (0, 255, 0), -1)
                    if(i < len(points) - 1):  # 最后一个点不需要画线段
                        cv2.line(frame, p, points[i+1], (255, 0, 0), 2)
                # ------------------------------------------------------

        cv2.imshow("det", frame)
        if( cv2.waitKey(1) == ord('q') ):
            break
        # vout.write(frame) # 保存为视频

    # vout.release() # 释放
